from gymnasium.envs.registration import register
from .make_panda_env import make_env
from .panda_env import PandaPickAndPlaceEnv,PandaPushEnv
from .utils import get_state,states_to_result

__all__ = ['make_env','PandaPickAndPlaceEnv','PandaPushEnv','get_state','states_to_result']

ENV_IDS = []

for task in ["Push", "PickAndPlace"]:
    for reward_type in ["sparse", "dense"]:
        for control_type in ["ee", "joints"]:
            for height in [0.2, 0.5, 1.0, 2.0]:
                for latent_dim in [5, 10]:
                    reward_suffix = "Dense" if reward_type == "dense" else ""
                    control_suffix = "Joints" if control_type == "joints" else ""
                    env_id = f"Panda{task}-{control_suffix}-{reward_suffix}-{height}-{latent_dim}-v3"

                    register(
                        id=env_id,
                        entry_point=f"stable_meta_learning.envs:Panda{task}Env",
                        kwargs={"reward_type": reward_type, "control_type": control_type, 'object_height':height, 'latent_dim': latent_dim},
                        max_episode_steps=100 if task == "Stack" else 50,
                    )

                    ENV_IDS.append(env_id)

